package top.likeqc.leetcode.editor.cn;

/**
 * @see <a href="https://leetcode.cn/problems/kth-ancestor-of-a-tree-node">1483. 树节点的第 K 个祖先</a>
 */
public class KthAncestorOfATreeNode {
    public static void main(String[] args) {
        TreeAncestor treeAncestor = new TreeAncestor(7, new int[] {-1, 0, 0, 1, 1, 2, 2});

        System.out.println(treeAncestor.getKthAncestor(3, 1));
        System.out.println(treeAncestor.getKthAncestor(5, 2));
        System.out.println(treeAncestor.getKthAncestor(6, 3));
    }
    // leetcode submit region begin(Prohibit modification and deletion)
    static class TreeAncestor {
        int[][] pa;

        public TreeAncestor(int n, int[] parent) {
            int m = 32 - Integer.numberOfLeadingZeros(n);
            pa = new int[n][m];

            for (int i = 0; i < n; i++) {
                pa[i][0] = parent[i];
            }

            for (int i = 0; i < m - 1; i++) {
                for (int j = 0; j < n; j++) {
                    int p = pa[j][i];
                    pa[j][i + 1] = p < 0 ? -1 : pa[p][i];
                }
            }
        }

        public int getKthAncestor(int node, int k) {
            int m = 32 - Integer.numberOfLeadingZeros(k);

            for (int i = 0; i < m; i++) {
                if (((k >> i) & 1) > 0) {
                    node = pa[node][i];
                    if (node < 0) {
                        break;
                    }
                }
            }

            return node;
        }
    }

    /**
     * Your TreeAncestor object will be instantiated and called as such: TreeAncestor obj = new
     * TreeAncestor(n, parent); int param_1 = obj.getKthAncestor(node,k);
     */
    // leetcode submit region end(Prohibit modification and deletion)

}
